package _13_red_black_tree.case3;

import static _13_red_black_tree.case3.RedBlackTree.Color.BLACK;
import static _13_red_black_tree.case3.RedBlackTree.Color.RED;

/**
 * @author: mornd
 * @dateTime: 2023/6/16 - 15:17
 * 红黑树定义
 */
class RedBlackTree<E> {
    RBNode root;

    enum Color {
        RED, BLACK;
    }

    static class RBNode<E> {
        int key;
        E value;
        RBNode left;
        RBNode right;
        RBNode parent;
        Color color = RED; // 默认红色

        public RBNode(int key) {
            this.key = key;
        }

        public RBNode(int key, Color color, RBNode left, RBNode right) {
            this.key = key;
            this.left = left;
            this.right = right;
            this.color = color;
            if(left != null) {
                left.parent = this;
            }
            if(right != null) {
                right.parent = this;
            }
        }

        public RBNode(int key, Color color) {
            this.key = key;
            this.color = color;
        }

        RBNode(int key, E value) {
            this.key = key;
            this.value = value;
        }

        // 当前节点是否是父亲的左孩子
        boolean isLeftChild() {
            return parent != null && parent.left == this;
        }

        // 获取叔叔节点
        RBNode uncle() {
            if (parent == null || parent.parent == null) {
                return null;
            }
            if (parent.isLeftChild()) {
                return parent.parent.right;
            } else {
                return parent.parent.left;
            }
        }

        // 获取兄弟节点
        RBNode sibling() {
            if (parent == null) {
                return null;
            }
            if (isLeftChild()) {
                return parent.right;
            } else {
                return parent.left;
            }
        }
    }

    /**
     * 删除节点
     */
    public void remove(int key) {
        RBNode deleted = find(key);
        if (deleted == null) {
            return;
        }
        remove(deleted);
    }

    /**
     * 删除 对应6种情况，删除前树是平衡的，不可能有红红相邻
     * cas0：待删除节点有两个孩子，李代桃僵，简化成最多只有一个孩子
     * case1：删除节点是根节点（又分根没有孩子和根只有一个孩子，如果根有两个孩子又会简化成删后继节点）
     * 如果待删除的节点是红色（不管该红有没有孩子）不需要考虑平衡（直接走删除逻辑），否则需要考虑平衡，变色+旋转
     * cse2：删的是黑，剩下的是红，将剩下的红变黑即可
     *      此时已经触发双黑（删除的和剩下的都是黑色）
     *          case3：被调整节点兄弟为红，旋转 + 父亲变红，兄弟变黑，再转化为 case4 或 case5 这种情况
     *          case4：被调整节点兄弟为黑，两个侄子都为黑时，又对应两种情况（父亲是红，父亲是黑）
     *          case5：被调整节点兄弟为黑，至少一个侄子为红时，又对应4种失衡情况（LL，LR，RL，RR）
     *
     * 删的是黑，剩的是黑  双黑，
     * 删的是黑，剩的是红  剩的变黑即可
     * 删的是红，剩的是黑  不用处理
     * 删的是红，剩的是红  不会出现
     */
    private void remove(RBNode deleted) {
        // 替代节点
        RBNode replaced = findReplaced(deleted);
        RBNode parent = deleted.parent; // 如果 deleted != root 那么 parent != null 就成立

        // 情况1：没有孩子
        if (replaced == null) {
            if (root == deleted) {
                // case1：根置为 null
                root = null;
            } else {
                if (isBlack(deleted)) {
                    // 待删除节点是黑色时处理，待删除节点如果是红色，不用调整平衡，直接删
                    fixDoubleBlack(deleted);
                }

                // 待删除的不是根节点，并且没有孩子
                if (deleted.isLeftChild()) {
                    parent.left = null;
                } else {
                    parent.right = null;
                }
                deleted.parent = null; // help GC
            }
            return;
        }

        // 情况2：只有一个孩子
        if (deleted.left == null || deleted.right == null) {
            if (root == deleted) {
                // case1：拿到孩子的k、v，孩子替死
                root.key = replaced.key;
                root.value = replaced.value;
                root.left = root.right = null;
            } else {
                // 待删除的不是根节点，并且有一个孩子
                if (deleted.isLeftChild()) {
                    parent.left = replaced;
                } else {
                    parent.right = replaced;
                }
                replaced.parent = parent;
                deleted.left = deleted.right = deleted.parent = null; // help GC

                // 先删除后调整
                if (isBlack(deleted)) {
                    if (isBlack(replaced)) {
                        // 删除的和剩下的都是黑时
                        fixDoubleBlack(replaced);
                    } else {
                        // 删的是黑，剩的是红，将剩下的红变黑即可
                        replaced.color = BLACK;
                    }
                }

//                if(isBlack(deleted) && isBlack(replaced)) {
//                    // 删除的和剩下的都是黑时
//                    fixDoubleBlack(replaced);
//                } else {
//                    // 删的是黑，剩的是红（或删的是红，剩的是黑），将剩下的红变黑即可
//                    replaced.color = BLACK;
//                }
            }
            return;
        }

        /*
            情况3：有两个孩子
            此时的 replaced 就是后继节点
            交换 后继节点 的key、value，李代桃僵
            将 replaced 的属性赋值给 deleted，这样就不用删除 deleted，而是去删除 replaced
            由于后继节点最多就一个孩子，这样递归调用一次，两个孩子删除就简化成了一个孩子删除
         */

        // cas0：待删除节点有两个孩子，李代桃僵，简化成最多只有一个孩子
        int k = deleted.key;
        deleted.key = replaced.key;
        replaced.key = k;

        E v = (E) deleted.value;
        deleted.value = replaced.value;
        replaced.value = v;

        remove(replaced); // 替死
    }

    /**
     * 处理双黑情况（少了一个黑）
     * case3：被调整节点兄弟为红（两个侄子必定为黑），旋转 + 父亲变红，兄弟变黑，再转化为 case4 或 case5 这种情况
     * case4：被调整节点兄弟为黑，两个侄子都为黑时（两个侄子不一定都为黑），先将兄弟变红，然后又对应两种情况（父亲是红，父亲是黑）
     * case5：被调整节点兄弟为黑，至少一个侄子为红时，又对应4种失衡情况（LL，LR，RL，RR）
     */
    private void fixDoubleBlack(RBNode node) {
        if(node == root) {
            return;
        }
        RBNode parent = node.parent;
        RBNode sibling = node.sibling();

        if(isRed(sibling)) {
            // case3 兄弟为红色
            if(node.isLeftChild()) {
                // node 是左孩子，父亲左旋
                leftRotate(parent);
            } else {
                rightRotate(parent);
            }
            // 变色
            parent.color = RED;
            sibling.color = BLACK;

            // 此时已经转化为情况 case4 或 case5 了，递归一次
            fixDoubleBlack(node);
        } else {
            if(sibling != null) {
                if(isBlack(sibling.left) && isBlack(sibling.right)) {
                    // case4：兄弟和两个侄子都为黑
                /*
                    先将兄弟变红，在判断父亲颜色
                        1、父亲为红：将父亲变黑即可，避免红红，此时黑色节点数目保存平衡了
                        2、父亲为黑：递归，再次让父亲触发双黑
                 */
                    sibling.color = RED;
                    if(isRed(parent)) {
                        parent.color = BLACK;
                    } else {
                        fixDoubleBlack(parent);
                    }
                } else {
                    // case5：兄弟为黑色，有一个或两个侄子为红色
                    /*
                        case5又分4情况 LL LR RL RR
                        兄弟是左孩子，左侄子是红 LL不平衡，父亲右旋，兄弟改为父亲颜色，父亲和侄子变黑
                        兄弟是左孩子，右侄子是红 LR不平衡，兄弟先左旋，父亲再右旋，此时侄子已替代父亲位置，侄子改为父亲颜色，父亲改为黑色
                        兄弟是右孩子，右侄子是红 RR不平衡，同上处理
                        兄弟是右孩子，左侄子是红 RL不平衡，
                     */
                    if(sibling.isLeftChild() && isRed(sibling.left)) {
                        // LL
                        rightRotate(parent);
                        sibling.color = parent.color;
                        sibling.left.color = BLACK; // 即使有红右侄子也不影响平衡
                    } else if (sibling.isLeftChild() && isRed(sibling.right)) {
                        sibling.right.color = parent.color;
                        // LR
                        leftRotate(sibling);
                        rightRotate(parent);
                    } else if (!sibling.isLeftChild() && isRed(sibling.left)) {
                        //RL
                        sibling.left.color = parent.color;
                        rightRotate(sibling);
                        leftRotate(parent);
                    } else {
                        // RR
                        leftRotate(parent);
                        sibling.color = parent.color;
                        sibling.right.color = BLACK;
                    }
                    parent.color = BLACK;
                }
            } else {
                fixDoubleBlack(parent);
            }
        }
    }

    /**
     * 返回待删除节点的剩余节点
     * 就是说我要删除节点a，那么谁来代替a呢，该方法的返回值就是来替代a的
     */
    private RBNode findReplaced(RBNode deleted) {
        if (deleted.left == null && deleted.right == null) {
            // 情况1：节点没有孩子，返回 null
            return null;
        }

        // 情况2：有一个孩子，返回那个孩子
        if (deleted.left == null) {
            return deleted.right;
        }
        if (deleted.right == null) {
            return deleted.left;
        }

        // 情况3：有两个孩子，返回后继节点
        RBNode s = deleted.right;
        while (s.left != null) {
            s = s.left;
        }
        return s;
    }

    /**
     * 根据 key 查找节点
     */
    private RBNode find(int k) {
        RBNode curr = root;
        while (curr != null) {
            if (k < curr.key) {
                curr = curr.left;
            } else if (k > curr.key) {
                curr = curr.right;
            } else {
                return curr;
            }
        }
        return null;
    }

    /**
     * 根据 key 查找节点值
     */
    public E get(int key) {
        RBNode node = find(key);
        return node == null ? null : (E) node.value;
    }

    /**
     * 新增，遇到红红不平衡（4种情况）进行调整，否则正常新增
     */
    public void put(int key, E value) {
        RBNode curr = root;
        RBNode parent = null;
        while (curr != null) {
            parent = curr;
            if (key < curr.key) {
                curr = curr.left;
            } else if (key > curr.key) {
                curr = curr.right;
            } else {
                curr.value = value;
                return;
            }
        }
        RBNode inserted = new RBNode(key, value);
        if (parent == null) {
            root = inserted;
        } else if (key < parent.key) {
            parent.left = inserted;
            inserted.parent = parent;
        } else {
            parent.right = inserted;
            inserted.parent = parent;
        }

        fixRedRed(inserted);
//        System.out.println(this + "\n-----------------------------------------");
    }

    /**
     * 处理红红的4种情况
     * case1：插入节点就是根节点，将根节点变黑即可
     * case2：新插入的节点默认是红色，此时只要父亲节点是黑色就无需调整
     * case3：uncle 为红色，处理：将父亲和叔叔变黑，爷爷变红
     * case4：当前节点node和父亲节点都是红，叔叔是黑或者是空，又分4种失衡情况 LL，LR，RL，RR
     */
    private void fixRedRed(RBNode node) {
        // case1：插入节点就是根节点，将根节点变黑即可
        if (root == node) {
            node.color = BLACK;
            return;
        }

        if (isBlack(node.parent)) {
            // case2：新插入的节点默认是红色，此时只要父亲节点是黑色就无需调整
            return;
        }

        // 此时红红相邻了
        RBNode parent = node.parent;
        RBNode uncle = node.uncle();
        RBNode grandparent = parent.parent;

        if (isRed(uncle)) {
            // 能进这里表示 uncle 不为空并且是红色，并且 node 和 node.parent 也都是红色
            // case3：uncle 为红色，处理：将父亲和叔叔变黑，爷爷变红
            parent.color = BLACK;
            uncle.color = BLACK;
            grandparent.color = RED;

            // 递归处理爷爷节点，因为多了黑色节点，其它地方又不平衡了，并且还要保证根为黑色
            fixRedRed(grandparent);
        } else {
            // case4：当前节点node和父亲节点都是红，叔叔是黑或者是空，又分4种失衡情况 LL，LR，RL，RR
            if (parent.isLeftChild() && node.isLeftChild()) {
                // true && true
                /*
                    LL 左左不平衡
                    父亲是左孩子，node（或者叫插入，因为会递归调用）节点是左孩子
                    处理：父亲变黑，爷爷变红，爷爷右旋
                 */
                parent.color = BLACK;
                grandparent.color = RED;
                rightRotate(grandparent);
            } else if (parent.isLeftChild()) {
                // true && false
                // LR  父亲是左孩子，当前节点是右孩子
                /*
                    LR 左右不平衡
                    父亲是左孩子，插入节点是右孩子
                    处理：父亲先左旋，当前节点变黑，爷爷变红，爷爷再右旋
                 */
                leftRotate(parent);
                node.color = BLACK;
                grandparent.color = RED;
                rightRotate(grandparent);
            } else if (node.isLeftChild()) {
                // false && true
                // LR
                rightRotate(parent);
                node.color = BLACK;
                grandparent.color = RED;
                leftRotate(grandparent);
            } else {
                // false && false
                // RR
                parent.color = BLACK;
                grandparent.color = RED;
                leftRotate(grandparent);
            }
        }
    }

    // 右旋
    private void rightRotate(RBNode pink) {
        RBNode parent = pink.parent;

        // 开始右旋，参照图片 img/右旋前.png
        RBNode yellow = pink.left;
        RBNode green = yellow.right;
        yellow.right = pink;
        pink.left = green;

        // 3个节点全部换爹
        yellow.parent = parent;
        pink.parent = yellow;
        if (green != null) {
            green.parent = pink;
        }

        // // 建立父节点与右旋节点的关系
        if (parent != null) {
            if (parent.left == pink) {
                parent.left = yellow;
            } else {
                parent.right = yellow;
            }
        } else {
            root = yellow;
        }
    }

    private void leftRotate(RBNode pink) {
        RBNode parent = pink.parent;

        RBNode yellow = pink.right;
        RBNode green = yellow.left;
        yellow.left = pink;
        pink.right = green;

        yellow.parent = parent;
        pink.parent = yellow;
        if (green != null) {
            green.parent = pink;
        }

        if (parent != null) {
            if (parent.left == pink) {
                parent.left = yellow;
            } else {
                parent.right = yellow;
            }
        } else {
            root = yellow;
        }
    }

    // 判断是否是红色 如果节点为 null，则为黑色
    private boolean isRed(RBNode n) {
        return n != null && n.color == RED;
    }

    private boolean isBlack(RBNode n) {
        return !isRed(n);
    }

    public boolean contains(int key) {
        return find(key) != null;
    }

    @Override
    public String toString() {
        StringBuilder sb = new StringBuilder();
        printVerticalTree(root, "", sb);
        return sb.toString();
    }

    private void printVerticalTree(RBNode node, String prefix, StringBuilder sb) {
        if (node == null) {
            return;
        }
        printVerticalTree(node.right, prefix + "        ", sb);
        sb.append(node.color == RED ? "\u001B[31m" : "\u001B[30m");
        sb.append(prefix).append("|--");
        sb.append(node.key).append("\u001B[0m").append("\n");
        printVerticalTree(node.left, prefix + "        ", sb);
    }
}
